#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
公共装饰器
"""

from functools import wraps
from flask import request, jsonify
import jwt

def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.headers.get('Authorization')
        if not token:
            return jsonify({'message': '缺少访问令牌'}), 401
        
        try:
            if token.startswith('Bearer '):
                token = token[7:]
            data = jwt.decode(token, 'billiard-jwt-secret', algorithms=['HS256'])
            from models.user import User
            current_user = User.query.get(data['user_id'])
            if not current_user:
                return jsonify({'message': '无效的令牌'}), 401
        except jwt.ExpiredSignatureError:
            return jsonify({'message': '令牌已过期'}), 401
        except jwt.InvalidTokenError:
            return jsonify({'message': '无效的令牌'}), 401
        
        return f(current_user, *args, **kwargs)
    return decorated

def admin_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.headers.get('Authorization')
        if not token:
            return jsonify({'message': '缺少访问令牌'}), 401
        
        try:
            if token.startswith('Bearer '):
                token = token[7:]
            data = jwt.decode(token, 'billiard-jwt-secret', algorithms=['HS256'])
            from models.admin import Admin
            admin = Admin.query.get(data['admin_id'])
            if not admin:
                return jsonify({'message': '无管理员权限'}), 403
        except jwt.ExpiredSignatureError:
            return jsonify({'message': '令牌已过期'}), 401
        except jwt.InvalidTokenError:
            return jsonify({'message': '无效的令牌'}), 401
        
        return f(admin, *args, **kwargs)
    return decorated